"""
任务管理器
"""

import asyncio
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any

from fastapi import HTTPException
from crawlee.sessions import Session

from models.api_models import TaskRequest, TaskResult


class TaskManager:
    """任务管理器"""
    
    def __init__(self, session_manager, browser_manager):
        self.session_manager = session_manager
        self.browser_manager = browser_manager
        self.tasks: Dict[str, TaskResult] = {}
        self.running_tasks: Dict[str, asyncio.Task] = {}
    
    async def execute_task(self, task_request: TaskRequest) -> str:
        """执行任务"""
        task_id = str(uuid.uuid4())
        
        # 创建任务结果
        task_result = TaskResult(
            task_id=task_id,
            url=task_request.url,
            session_id=task_request.session_id,
            browser_id=task_request.browser_id,
            status="pending",
            result=None,
            error=None,
            created_at=datetime.now(),
            completed_at=None
        )
        
        self.tasks[task_id] = task_result
        
        # 启动后台任务
        background_task = asyncio.create_task(
            self._run_task(task_id, task_request)
        )
        self.running_tasks[task_id] = background_task
        
        return task_id
    
    async def _run_task(self, task_id: str, task_request: TaskRequest):
        """运行具体任务"""
        try:
            self.tasks[task_id].status = "running"
            
            # 获取Session
            session = None
            if task_request.session_id:
                session = self.session_manager.get_session(task_request.session_id)
                if not session:
                    raise HTTPException(status_code=404, detail=f"Session {task_request.session_id} not found")
            
            # 获取浏览器
            browser_data = None
            if task_request.browser_id:
                browser_data = self.browser_manager.get_browser(task_request.browser_id)
                if not browser_data:
                    raise HTTPException(status_code=404, detail=f"Browser {task_request.browser_id} not found")
            
            # 执行爬取任务
            result = await self._crawl_with_session_browser(
                task_request.url,
                session,
                browser_data,
                task_request.handler_type,
                task_request.custom_data or {}
            )
            
            # 更新任务结果
            self.tasks[task_id].status = "completed"
            self.tasks[task_id].result = result
            self.tasks[task_id].completed_at = datetime.now()
            
        except Exception as e:
            self.tasks[task_id].status = "failed"
            self.tasks[task_id].error = str(e)
            self.tasks[task_id].completed_at = datetime.now()
        
        finally:
            # 清理运行中的任务
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
    
    async def _crawl_with_session_browser(self, 
                                        url: str, 
                                        session: Optional[Session],
                                        browser_data: Optional[Dict[str, Any]],
                                        handler_type: str,
                                        custom_data: Dict[str, Any]) -> Dict[str, Any]:
        """使用指定Session和浏览器执行爬取"""
        
        if browser_data:
            # 使用指定浏览器
            context = browser_data["context"]
            page = await context.new_page()
            browser_data["last_used"] = datetime.now()
        else:
            # 创建临时浏览器
            from playwright.async_api import async_playwright
            playwright = await async_playwright().start()
            browser = await playwright.chromium.launch(headless=True)
            context = await browser.new_context()
            page = await context.new_page()
        
        try:
            # 访问页面
            await page.goto(url)
            
            # 获取cookies
            browser_cookies = await context.cookies()
            
            # 更新Session中的cookie信息
            if session:
                session.user_data['custom_cookies'] = {
                    cookie['name']: cookie['value'] for cookie in browser_cookies
                }
                session.user_data['last_activity'] = datetime.now()
                session.user_data['last_url'] = url
            
            # 根据handler类型执行不同逻辑
            if handler_type == "login":
                result = await self._handle_login(page, session, custom_data)
            elif handler_type == "custom":
                result = await self._handle_custom(page, session, custom_data)
            else:
                result = await self._handle_default(page, session, custom_data)
            
            # 标记session状态
            if session:
                if result.get("status") == "success":
                    session.mark_good()
                else:
                    session.mark_bad()
            
            return result
            
        finally:
            if not browser_data:
                # 关闭临时浏览器
                await context.close()
                await browser.close()
                await playwright.stop()
            else:
                # 只关闭页面
                await page.close()
    
    async def _handle_default(self, page, session: Optional[Session], custom_data: Dict[str, Any]) -> Dict[str, Any]:
        """默认处理器"""
        title = await page.title()
        content_preview = (await page.content())[:500]
        url = page.url
        
        return {
            "handler_type": "default",
            "url": url,
            "title": title,
            "content_preview": content_preview,
            "session_id": session.id if session else None,
            "cookies_count": len(session.user_data.get('custom_cookies', {})) if session else 0,
            "status": "success",
            "timestamp": datetime.now().isoformat()
        }
    
    async def _handle_login(self, page, session: Optional[Session], custom_data: Dict[str, Any]) -> Dict[str, Any]:
        """登录处理器"""
        # 可以在这里实现具体的登录逻辑
        username = custom_data.get("username")
        password = custom_data.get("password")
        
        login_result = {
            "handler_type": "login",
            "url": page.url,
            "login_attempted": bool(username and password),
            "session_id": session.id if session else None,
            "status": "success"
        }
        
        if username and password:
            # 这里可以添加具体的登录操作
            # 例如：填写表单、点击登录按钮等
            login_result["message"] = f"Login attempted for user: {username}"
        
        return login_result
    
    async def _handle_custom(self, page, session: Optional[Session], custom_data: Dict[str, Any]) -> Dict[str, Any]:
        """自定义处理器"""
        # 用户可以在custom_data中传入自定义的处理逻辑
        selectors = custom_data.get("selectors", {})
        actions = custom_data.get("actions", [])
        
        result = {
            "handler_type": "custom",
            "url": page.url,
            "session_id": session.id if session else None,
            "extracted_data": {},
            "actions_performed": [],
            "status": "success"
        }
        
        # 执行数据提取
        for name, selector in selectors.items():
            try:
                element = await page.wait_for_selector(selector, timeout=5000)
                if element:
                    text = await element.text_content()
                    result["extracted_data"][name] = text
            except:
                result["extracted_data"][name] = None
        
        # 执行操作
        for action in actions:
            try:
                action_type = action.get("type")
                if action_type == "click":
                    await page.click(action["selector"])
                    result["actions_performed"].append(f"Clicked: {action['selector']}")
                elif action_type == "fill":
                    await page.fill(action["selector"], action["value"])
                    result["actions_performed"].append(f"Filled: {action['selector']}")
                elif action_type == "wait":
                    await page.wait_for_timeout(action.get("timeout", 1000))
                    result["actions_performed"].append(f"Waited: {action.get('timeout', 1000)}ms")
            except Exception as e:
                result["actions_performed"].append(f"Failed: {action} - {str(e)}")
        
        return result
    
    def get_task(self, task_id: str) -> Optional[TaskResult]:
        """获取任务结果"""
        return self.tasks.get(task_id)
    
    def get_all_tasks(self) -> List[TaskResult]:
        """获取所有任务"""
        return list(self.tasks.values())
    
    def get_running_tasks(self) -> List[str]:
        """获取运行中的任务ID"""
        return list(self.running_tasks.keys()) 